#generate_split
import os,json
import random
def generate_predict_split(data_path):
    '''
    all_split generate to all.json
    '''
    with open(os.path.join(data_path,'annotations.json')) as f:
        data_dict=json.load(f)
    all_split={'train':[],'test':[]}
    for image_name in data_dict:
        all_split['test'].append(image_name)
    with open(os.path.join(data_path,'split','all.json'),'w') as f:
        json.dump(all_split,f)
def get_indepent_split(data_path):
    '''
    Generate three splits where each part takes turns being the validation set.
    Save to val_1.json, val_2.json, and val_3.json.
    Test split remains empty in all cases.
    '''
    # 读取 annotations.json
    with open(os.path.join(data_path, 'annotations.json'), 'r') as f:
        data_dict = json.load(f)
    
    # 筛选包含 diagnosis 和 text 的图像
    label_image_list = [
        image_name for image_name in data_dict 
        if 'diagnosis' in data_dict[image_name] and 'text' in data_dict[image_name]['diagnosis']
    ]
    
    # 检查数据是否足够分成三份
    if len(label_image_list) < 3:
        raise ValueError("Not enough images to split into three parts.")
    
    # 随机打乱图像列表
    random.shuffle(label_image_list)
    
    # 计算每份的大小（尽量均匀）
    total_images = len(label_image_list)
    part_size = total_images // 3  # 每份的基本大小
    remainder = total_images % 3   # 余数，用于分配到前几份
    
    # 分成三份
    parts = []
    start = 0
    for i in range(3):
        # 如果有余数，前几份多分配一个
        size = part_size + (1 if i < remainder else 0)
        parts.append(label_image_list[start:start + size])
        start += size
    
    # 轮流生成三个分割
    for val_idx in range(3):
        split = {'train': [], 'val': [], 'test': []}
        
        # 当前部分作为 val
        split['val'] = parts[val_idx]
        
        # 其余部分作为 train
        split['train'] = parts[(val_idx + 1) % 3] + parts[(val_idx + 2) % 3]
        
        # test 保持为空
        split['test'] = []
        
        # 保存到文件
        output_file = os.path.join(data_path, 'split',f'val_{val_idx + 1}.json')
        with open(output_file, 'w') as f:
            json.dump(split, f, indent=4)
        
        print(f"Saved split {val_idx + 1} to {output_file}")
    
    return  # 不返回具体分割，因为结果已保存到文件中
if __name__ == "__main__":
    get_indepent_split(data_path="/mnt/c/DocumentWorkSpace/public_processes/APTOS")